/**
 * 
 */
package edu.umd.clip.lm.model.data;

import java.io.IOException;
import java.util.*;

import edu.umd.clip.lm.factors.FactorTuple;
import edu.umd.clip.lm.util.LRU;
import edu.umd.clip.lm.util.Long2IntMap;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class ContextReducingReadableTrainingData extends ReadableTrainingData {
	private final byte newContextSize;
	private final int numNullOvertFactors;
	private final int numNullHiddenFactors;
	private final Cache cache;
	/**
	 * @param reader
	 */
	public ContextReducingReadableTrainingData(TrainingDataReader reader, int overtOrder, int hiddenOrder, int cacheSize) {
		super(reader);
		this.newContextSize = (byte) Math.max(overtOrder-1, hiddenOrder-1);
		
		this.numNullOvertFactors = Math.min(newContextSize + 1 - overtOrder, newContextSize);
		this.numNullHiddenFactors = Math.min(newContextSize + 1 - hiddenOrder, newContextSize);
		
		this.cache = new Cache(cacheSize);
		super.setContextSize(newContextSize);
	}
	
	@SuppressWarnings("serial")
	private class Cache extends LRU<Context, Long2IntMap> {
		TrainingDataBlock cacheBlock;
		/**
		 * @param initialCapacity
		 */
		public Cache(int initialCapacity) {
			super(initialCapacity);
			cacheBlock = new TrainingDataBlock();
		}
		@Override
		protected boolean removeEldestEntry(Map.Entry<Context, Long2IntMap> eldest) {
			if (super.removeEldestEntry(eldest)) {
				ContextFuturesPair pair = new ContextFuturesPair(eldest.getKey(), TupleCountPair.fromMap(eldest.getValue()));
				cacheBlock.add(pair);
				return true;
			}
			return false;
		}
	}

	@Override
	public TrainingDataBlock next() throws IOException {
		TrainingDataBlock result = null;
		
		while(super.hasNext()) {
			TrainingDataBlock currentBlock = super.next();
			if (currentBlock == null) {
				// feed the cache
				break;
			}
			//System.err.printf("ContextReduce: reading block, totalCount = %d\n", currentBlock.getTotalCount());
			
			for(ContextFuturesPair pair : currentBlock) {
				Context ctx = reduceContext(pair.getContext());
				
				Long2IntMap futures = cache.get(ctx);
				if (futures == null) {
					futures = new Long2IntMap(pair.getFutures().length + 2);
					cache.put(ctx, futures);
				}
				for(TupleCountPair tc : pair.getFutures()) {
					futures.addAndGet(tc.tuple, tc.count);
				}

				if (cache.cacheBlock.size() >= TrainingDataBlock.getMaxCapacity()) {
					result = cache.cacheBlock;
					cache.cacheBlock = new TrainingDataBlock();
				}
			}
			if (result != null) return result;
		}
		// read everything, drain the cache now
		for(Iterator<Map.Entry<Context, Long2IntMap>> it = cache.entrySet().iterator(); it.hasNext();) {
			Map.Entry<Context, Long2IntMap> entry = it.next();
			it.remove();
			
			Context ctx = entry.getKey();
			TupleCountPair futures[] = TupleCountPair.fromMap(entry.getValue());
			ContextFuturesPair pair = new ContextFuturesPair(ctx, futures);
			if (!cache.cacheBlock.add(pair)) {
				result = cache.cacheBlock;
				cache.cacheBlock = new TrainingDataBlock();
				cache.cacheBlock.add(pair);
				return result;
			}
		}
		if (cache.cacheBlock.hasData()) {
			result = cache.cacheBlock;
			cache.cacheBlock = new TrainingDataBlock();
			return result;
		}
		return null;
	}

	private Context reduceContext(Context oldContext) {
		final int oldContextSize = oldContext.data.length;

		Context ctx = new Context(newContextSize+1);
		
		if (newContextSize == oldContextSize) {
			ctx.data = Arrays.copyOf(oldContext.data, oldContext.data.length); 
		} else {
			for(int i=oldContextSize-newContextSize; i<oldContextSize; ++i) {
				ctx.data[i - oldContextSize+newContextSize] = oldContext.data[i];
			}
		}
		for(int i=0; i<numNullOvertFactors; ++i) {
			ctx.data[i] &= FactorTuple.getHiddenMask();
		}
		for(int i=0; i<numNullHiddenFactors; ++i) {
			ctx.data[i] &= FactorTuple.getOvertMask();
		}
		return ctx;
	}

	@Override
	public boolean hasNext() {
		return !cache.isEmpty() || cache.cacheBlock.hasData() || super.hasNext();
	}

	@Override
	public byte getContextSize() {
		return newContextSize;
	}

	@Override
	public void setContextSize(byte contextSize) {
	}
}
